"""
构建属性训练集，也就是输入一句话，预定好的属性，得出属性的值
得到attributes.json这个训练集合，这个训练集主要是问某个描述的属性的
"""
import copy
import json
import os

import pandas as pd

from extract_attributes.main import read_pdf_from_file, get_similarity, get_allFile_from_folder, M_jaccard
from pdf2excelv1.main import text2sentence


def out2file(filename, data):
    with open(filename, "w+", encoding="utf-8") as f:
        f.write(json.dumps(data, ensure_ascii=False, indent=2) + "\n")


"""
获得星形结构，也就是entity与property的关联结构

流程
1、通过entity的值进行分割
2、如果property=描述,value文本；否则property=其他内容，value为属性

输入：excel标注文档
输出：星型结构
"""


def get_star_structure(file_path):
    df = pd.read_excel(file_path, engine='openpyxl')
    entitys = df['entity']
    propertys = df['property']
    values = df['value']
    dic = {}
    # 分组
    for i, entity in enumerate(entitys):
        # 找到所有的临床问题
        if entity not in dic:
            dic[entity] = {propertys[i]: values[i]}
        else:
            if propertys[i] not in dic[entity]:
                dic[entity].update({propertys[i]: values[i]})
            else:
                dic[entity].update({propertys[i]: str(dic[entity][propertys[i]]) + "\n" + str(values[i])})

    # 分组结束后，选出包含property=描述的dict
    dic2 = {}
    for key, value in dic.items():
        if "描述" in value:
            dic2.update({key: value})
    return dic2


"""
替换星型结构的描述字段，描述字段作为问题的主体，应该取自原文
通过相似度匹配（get_similarity），匹配到pdf中的文本内容
与extract_attributes中不同的是，先通过分句再匹配
"""


def replace_from_text(pdf_file_path, dic):
    # 通过文件名得到pdf的内容
    pdf_text = read_pdf_from_file(pdf_file_path).replace("\n", "").encode('gb2312', 'ignore').decode('gb2312')
    # 分句
    pdf_sentences = text2sentence(pdf_text)

    dic2 = copy.deepcopy(dic)
    for key, value in dic.items():
        sentence = value["描述"]
        max_jac = 0
        max_jac_pdf_sentence = ""
        for pdf_sentence in pdf_sentences:
            # get_similarity返回最大相似度，以及描述字段替换后的值
            # 获得对pdf分句结果中与标注数据中描述的相似度最大的句子
            jaccard = M_jaccard(sentence, pdf_sentence)
            if jaccard > max_jac:
                max_jac = jaccard
                max_jac_pdf_sentence = pdf_sentence
        dic2[key]["描述"] = max_jac_pdf_sentence
        if max_jac < 0.4:  # 相似度太低的样本就丢弃了，这个值很重要，要调节到一个合适的值
            dic2.pop(key)
    return dic2


"""
生成训练文本格式，{ask="",answer=""}
输入：星型结构
输出：训练格式的文本
"""


def get_train_format(dic):
    result = []
    for key, value in dic.items():
        text = value["描述"]
        for key2, value2 in value.items():
            if key2 == "描述":
                continue
            tep_dic = {"ask": "根据文本内容:“" + str(text) + "”回答" + str(key2) + "是什么", "answer": str(value2)}
            result.append(tep_dic)
    return result


"""
获得所有文档的训练格式的文本
"""


def get_all_train(folder_path=r'指南知识图谱三元组提取'):
    # 获得某个文件夹下的所有文件的路径，文件类型是file_type
    pdf_file_paths = (get_allFile_from_folder(folder_path, file_type="pdf"))
    excel_file_paths = (get_allFile_from_folder(folder_path, file_type="xlsx"))
    results = []
    for pdf_file_path in pdf_file_paths:
        dirname = os.path.dirname(pdf_file_path)
        # 找到与某个pdf同目录下的excel
        excel_file_path = [i for i in excel_file_paths if os.path.dirname(i) == dirname][0]
        # 星形结构，可见extract_attributes/example003.json
        orgin_dic = get_star_structure(excel_file_path)
        # 替换描述列
        replace_dic = replace_from_text(pdf_file_path, orgin_dic)
        # 替换前和替换后的都要，如果不想要替换前的，也可以删除
        results = results + get_train_format(orgin_dic) + get_train_format(replace_dic)
        # 这一步主要是把一些引用删除，可以写的更加严谨一点
        results = [x for x in results if not x["answer"].startswith("202")]
    return results


if __name__ == '__main__':
    out2file("attributes_traindata1.json", get_all_train(folder_path=r'../指南知识图谱三元组提取'))